using DataFrames
using QuadGK
using Distributions
using CSV
using LinearAlgebra
using JLD
using Interpolations
using MatrixEquations
using SparseArrays
using DelimitedFiles

#clearconsole()

#-------------------------------------------------------------
# IRF Configuration
#-------------------------------------------------------------

H  = 40 # maximum horizon
n_every = 10 # use every n_every'th draw from the posterior
sh_size = 3   # shock size, in multiples of standard deviations

#-------------------------------------------------------------
# include Functions
#-------------------------------------------------------------
#cd("$(pwd())/Dropbox/Heterogeneity/Software/KS_Simulation/")
readDir = "$(pwd())/Functions/"
#include(readDir *"vech.jl");
include(readDir *"logSpline_Procedures.jl");
#include(readDir *"VAR_Procedures.jl");
include(readDir *"Loaddata.jl");
include(readDir *"VAR_MoreLags_Procedures.jl")
include(readDir *"GiniFracBelowCutoff_Procedures.jl")
#include(readDir *"IRF_Procedures.jl")
include(readDir *"EmpPercentiles_Procedures.jl")


#-------------------------------------------------------------
# choose specification files
#-------------------------------------------------------------
nfVARSpec = "10tc"
nModSpec  = "1"
nMCMCSpec = "1"
modName   = "SS"  # VAR or SS

specDir   = "$(pwd())/SpecFiles/"
include(specDir * "/fVARspec" * nfVARSpec * ".jl")
include(specDir * "/" * modName * "spec" * nModSpec * ".jl")
include(specDir * "/" * modName * "MCMCspec" * nMCMCSpec * ".jl")

#-------------------------------------------------------------
# load aggregate data
#-------------------------------------------------------------
agg_data, period_agg, mean_unrate = loadaggdata(SampleStart,SampleEnd)#,juliaversion)
n_agg = size(agg_data)[2]

#-------------------------------------------------------------
# Load coefficients from density estimation
#-------------------------------------------------------------
sNameLoadDir = "fVAR" * nfVARSpec
loaddir  = "$(pwd())/results/" * sNameLoadDir *"/";
knots_all = readdlm(loaddir * "fVAR" *nfVARSpec * "_knots_all.csv", ',', Float64, '\n'; skipstart = 1);
ii=getindex.(findall(K_vec.-K.==0),[1 2])[1] # find index ii where K==K_vec
knots = knots_all[quant_sel[ii,:].==1]

PhatDensCoef_factor, MDD_GoF, VinvLam_all, period_Dens_ind, PhatDensCoef_lambda, PhatDensCoef_mean, PhatDensCoef_mean_allt  = loaddensdata(SampleStart,SampleEnd,K,nfVARSpec)
n_cross = size(PhatDensCoef_factor)[2]

#-------------------------------------------------------------
# Load posterior draws
#-------------------------------------------------------------

sName    = "fVAR" * nfVARSpec * "_" * modName * nModSpec * "_" * "MCMC" * nMCMCSpec;
loadDir  = "$(pwd())/results/" * sName *"/";

Bpmean     = load(loadDir * sName * "_PostMeans.jld", "Bpmean")
Sigpmean   = load(loadDir * sName * "_PostMeans.jld", "Sigpmean")
Sigtrpmean = load(loadDir * sName * "_PostMeans.jld", "Sigtrpmean")

Bpdraw     = load(loadDir * sName * "_PostDraws.jld", "Bpdraw")
Sigpdraw   = load(loadDir * sName * "_PostDraws.jld", "Sigpdraw")
Sigtrpdraw = load(loadDir * sName * "_PostDraws.jld", "Sigtrpdraw")

###################################################################
# find the optimal q's for the aggregate target variable
###################################################################

#-------------------------------------------------------------
# Generate IRFs at posterior mean of (Phi,Sigmatr)
# hh=1 is steady state
# hh=2 is period of impact
#-------------------------------------------------------------
println("")
println("Generating IRFs at posterior mean... ")
println("")

n               = n_agg+n_cross
p               = nlags

select_agg = 3 # select id of targeting variable (2: GDP, 3: employment)

YY_IRF          = zeros(H+1,n)
L               = Sigtrpmean
Btilde          = reshape(Bpmean, n*p,n)
response        = IRFredu_Anatomy_Agg(Btilde,L,H,sh_size,select_agg)
YY_IRF[2:end,:] = response'

PhatDens_IRF    = zeros(H+1,length(xgrid));
# compute density IRFs
for hh = 1:H+1
    PhatDensCoef_hh    = coefRecover(YY_IRF[hh,n_agg+1:end]', PhatDensCoef_lambda, PhatDensCoef_mean)
    #PhatDensNorm_hh    = lnpdfNormalize(PhatDensCoef_hh,knots,minimum(xgrid), maximum(xgrid))
    PhatDensNorm_hh    = lnpdfNormalize_unrate(PhatDensCoef_hh,knots,YY_IRF[hh,3]+mean_unrate,minimum(xgrid), maximum(xgrid))
    PhatDens_IRF[hh,:] = pdfEval(xgrid,PhatDensCoef_hh',knots,[PhatDensNorm_hh[1]]);
end

savedir = "$(pwd())/results_BCAnatomy/" * sName *"/";
try mkdir(savedir) catch; end
CSV.write(savedir * sName * "_IRF_PhatDens_AnatomySh_Agg" * string(select_agg) * "_pmean.csv", DataFrame(PhatDens_IRF))
CSV.write(savedir * sName * "_IRF_YY_AnatomySh_Agg" * string(select_agg) * "_pmean.csv", DataFrame(YY_IRF))

#-------------------------------------------------------------
# Generate IRFs for a subset of posterior draws
#-------------------------------------------------------------
println("")
println("Generating IRFs for subset of posterior draws... ")
println("")

nsim = size(Bpdraw)[1]
n_subseq                 = floor(Int,nsim/n_every)
YY_IRF_uncertainty       = zeros(H+1,n,n_subseq)

# hh=1 is steady state
# hh=2 is period of impact

for pp = 1:n_subseq

    time_init_loop = time_ns();
    println(" Draw number:  $pp")
    println(" Remaining draws:  $(n_subseq-pp)")
    println(" Posterior draw number: $(pp*n_every)")

    L      = Sigtrpdraw[pp*n_every,:,:]
    Btilde = Bpdraw[pp*n_every,:,:]
    response = IRFredu_Anatomy_Agg(Btilde,L,H,sh_size,select_agg)
    YY_IRF_uncertainty[2:end,:,pp] = response'

    PhatDens_IRF   = zeros(H+1,length(xgrid));

    # compute density IRFs
    for hh = 1:H+1
        PhatDensCoef_hh    = coefRecover(YY_IRF_uncertainty[hh,n_agg+1:end,pp]', PhatDensCoef_lambda, PhatDensCoef_mean)
    #    PhatDensNorm_hh    = lnpdfNormalize(PhatDensCoef_hh,knots, minimum(xgrid), maximum(xgrid))
        PhatDensNorm_hh    = lnpdfNormalize_unrate(PhatDensCoef_hh,knots,YY_IRF_uncertainty[hh,3,pp]+mean_unrate,minimum(xgrid), maximum(xgrid))
        PhatDens_IRF[hh,:] = pdfEval(xgrid,PhatDensCoef_hh',knots,[PhatDensNorm_hh[1]]);
    end

    CSV.write(savedir * sName * "_IRF_YY_AnatomySh_Agg" * string(select_agg) * "_" * string(pp) * ".csv", DataFrame(YY_IRF_uncertainty[:,:,pp]))
    CSV.write(savedir * sName * "_IRF_PhatDens_AnatomySh_Agg" * string(select_agg) * "_" * string(pp) * ".csv", DataFrame(PhatDens_IRF))

    time_loop=signed(time_ns()-time_init_loop)/1000000000
    println("Elapsed time = $(time_loop)")
    println("")

end

###################################################################
# find the optimal q's for the distribution (Method 1 in the note)
###################################################################

#-------------------------------------------------------------
# Generate IRFs at posterior mean of (Phi,Sigmatr)
# hh=1 is steady state
# hh=2 is period of impact
#-------------------------------------------------------------
println("")
println("Generating IRFs at posterior mean... ")
println("")

delta_x = xgrid[2]-xgrid[1]
ZZ = basis_logspline(xgrid, knots)

YY_IRF          = zeros(H+1,n)
L               = Sigtrpmean
Btilde          = reshape(Bpmean, n*p,n)
response        = IRFredu_Anatomy_Distr1(Btilde,L,H,sh_size)
YY_IRF[2:end,:] = response'

PhatDens_IRF    = zeros(H+1,length(xgrid));
# compute density IRFs
for hh = 1:H+1
    PhatDensCoef_hh    = coefRecover(YY_IRF[hh,n_agg+1:end]', PhatDensCoef_lambda, PhatDensCoef_mean)
    #PhatDensNorm_hh    = lnpdfNormalize(PhatDensCoef_hh,knots,minimum(xgrid), maximum(xgrid))
    PhatDensNorm_hh    = lnpdfNormalize_unrate(PhatDensCoef_hh,knots,YY_IRF[hh,3]+mean_unrate,minimum(xgrid), maximum(xgrid))
    PhatDens_IRF[hh,:] = pdfEval(xgrid,PhatDensCoef_hh',knots,[PhatDensNorm_hh[1]]);
end

savedir = "$(pwd())/results_BCAnatomy/" * sName *"/";
try mkdir(savedir) catch; end
CSV.write(savedir * sName * "_IRF_PhatDens_AnatomySh_Distr1_pmean.csv", DataFrame(PhatDens_IRF)) # method 1
CSV.write(savedir * sName * "_IRF_YY_AnatomySh_Distr1_pmean.csv", DataFrame(YY_IRF))

#-------------------------------------------------------------
# Generate IRFs for a subset of posterior draws
#-------------------------------------------------------------
println("")
println("Generating IRFs for subset of posterior draws... ")
println("")

nsim = size(Bpdraw)[1]
n_subseq                 = floor(Int,nsim/n_every)
YY_IRF_uncertainty       = zeros(H+1,n,n_subseq)

# hh=1 is steady state
# hh=2 is period of impact

for pp = 1:n_subseq

    time_init_loop = time_ns();
    println(" Draw number:  $pp")
    println(" Remaining draws:  $(n_subseq-pp)")
    println(" Posterior draw number: $(pp*n_every)")

    L      = Sigtrpdraw[pp*n_every,:,:]
    Btilde = Bpdraw[pp*n_every,:,:]
    response = IRFredu_Anatomy_Distr1(Btilde,L,H,sh_size)
    YY_IRF_uncertainty[2:end,:,pp] = response'

    PhatDens_IRF   = zeros(H+1,length(xgrid));

    # compute density IRFs
    for hh = 1:H+1
        PhatDensCoef_hh    = coefRecover(YY_IRF_uncertainty[hh,n_agg+1:end,pp]', PhatDensCoef_lambda, PhatDensCoef_mean)
    #    PhatDensNorm_hh    = lnpdfNormalize(PhatDensCoef_hh,knots, minimum(xgrid), maximum(xgrid))
        PhatDensNorm_hh    = lnpdfNormalize_unrate(PhatDensCoef_hh,knots,YY_IRF_uncertainty[hh,3,pp]+mean_unrate,minimum(xgrid), maximum(xgrid))
        PhatDens_IRF[hh,:] = pdfEval(xgrid,PhatDensCoef_hh',knots,[PhatDensNorm_hh[1]]);
    end

    CSV.write(savedir * sName * "_IRF_YY_AnatomySh_Distr1_" * string(pp) * ".csv", DataFrame(YY_IRF_uncertainty[:,:,pp]))
    CSV.write(savedir * sName * "_IRF_PhatDens_AnatomySh_Distr1_" * string(pp) * ".csv", DataFrame(PhatDens_IRF))

    time_loop=signed(time_ns()-time_init_loop)/1000000000
    println("Elapsed time = $(time_loop)")
    println("")

end

###################################################################
# find the optimal q's for the distribution (Method 2 in the note)
###################################################################
println("")
println("Generating IRFs at posterior mean... ")
println("")

YY_IRF          = zeros(H+1,n)
L               = Sigtrpmean
Btilde          = reshape(Bpmean, n*p,n)
response        = IRFredu_Anatomy_Distr2(Btilde,L,H,sh_size)
YY_IRF[2:end,:] = response'

PhatDens_IRF    = zeros(H+1,length(xgrid));
# compute density IRFs
for hh = 1:H+1
    PhatDensCoef_hh    = coefRecover(YY_IRF[hh,n_agg+1:end]', PhatDensCoef_lambda, PhatDensCoef_mean)
    #PhatDensNorm_hh    = lnpdfNormalize(PhatDensCoef_hh,knots,minimum(xgrid), maximum(xgrid))
    PhatDensNorm_hh    = lnpdfNormalize_unrate(PhatDensCoef_hh,knots,YY_IRF[hh,3]+mean_unrate,minimum(xgrid), maximum(xgrid))
    PhatDens_IRF[hh,:] = pdfEval(xgrid,PhatDensCoef_hh',knots,[PhatDensNorm_hh[1]]);
end

savedir = "$(pwd())/results_BCAnatomy/" * sName *"/";
try mkdir(savedir) catch; end
CSV.write(savedir * sName * "_IRF_PhatDens_AnatomySh_Distr2_pmean.csv", DataFrame(PhatDens_IRF)) # method 1
CSV.write(savedir * sName * "_IRF_YY_AnatomySh_Distr2_pmean.csv", DataFrame(YY_IRF))


#-------------------------------------------------------------
# Generate IRFs for a subset of posterior draws
#-------------------------------------------------------------
println("")
println("Generating IRFs for subset of posterior draws... ")
println("")

nsim = size(Bpdraw)[1]
n_subseq                 = floor(Int,nsim/n_every)
YY_IRF_uncertainty       = zeros(H+1,n,n_subseq)

# hh=1 is steady state
# hh=2 is period of impact

for pp = 1:n_subseq

    time_init_loop = time_ns();
    println(" Draw number:  $pp")
    println(" Remaining draws:  $(n_subseq-pp)")
    println(" Posterior draw number: $(pp*n_every)")

    L      = Sigtrpdraw[pp*n_every,:,:]
    Btilde = Bpdraw[pp*n_every,:,:]
    response = IRFredu_Anatomy_Distr2(Btilde,L,H,sh_size)
    YY_IRF_uncertainty[2:end,:,pp] = response'

    PhatDens_IRF   = zeros(H+1,length(xgrid));

    # compute density IRFs
    for hh = 1:H+1
        PhatDensCoef_hh    = coefRecover(YY_IRF_uncertainty[hh,n_agg+1:end,pp]', PhatDensCoef_lambda, PhatDensCoef_mean)
    #    PhatDensNorm_hh    = lnpdfNormalize(PhatDensCoef_hh,knots, minimum(xgrid), maximum(xgrid))
        PhatDensNorm_hh    = lnpdfNormalize_unrate(PhatDensCoef_hh,knots,YY_IRF_uncertainty[hh,3,pp]+mean_unrate,minimum(xgrid), maximum(xgrid))
        PhatDens_IRF[hh,:] = pdfEval(xgrid,PhatDensCoef_hh',knots,[PhatDensNorm_hh[1]]);
    end

    CSV.write(savedir * sName * "_IRF_YY_AnatomySh_Distr2_" * string(pp) * ".csv", DataFrame(YY_IRF_uncertainty[:,:,pp]))
    CSV.write(savedir * sName * "_IRF_PhatDens_AnatomySh_Distr2_" * string(pp) * ".csv", DataFrame(PhatDens_IRF))

    time_loop=signed(time_ns()-time_init_loop)/1000000000
    println("Elapsed time = $(time_loop)")
    println("")

end
